﻿#include <iostream>
#include <windows.h>
#include <mmdeviceapi.h>
#include <audioclient.h>
#include <functiondiscoverykeys_devpkey.h>
#include <fstream>
#include <vector>

#define REFTIMES_PER_SEC  10000000
#define REFTIMES_PER_MILLISEC  10000

const CLSID CLSID_MMDeviceEnumerator = __uuidof(MMDeviceEnumerator);
const IID IID_IMMDeviceEnumerator = __uuidof(IMMDeviceEnumerator);
const IID IID_IAudioClient = __uuidof(IAudioClient);
const IID IID_IAudioCaptureClient = __uuidof(IAudioCaptureClient);

HRESULT RecordAudioStream(const WCHAR* outputFile) {
    HRESULT hr;
    REFERENCE_TIME hnsRequestedDuration = REFTIMES_PER_SEC;
    REFERENCE_TIME hnsActualDuration;
    UINT32 bufferFrameCount;
    UINT32 numFramesAvailable;
    IMMDeviceEnumerator *pEnumerator = NULL;
    IMMDevice *pDevice = NULL;
    IAudioClient *pAudioClient = NULL;
    IAudioCaptureClient *pCaptureClient = NULL;
    WAVEFORMATEX *pwfx = NULL;
    UINT32 packetLength = 0;
    BYTE *pData;
    DWORD flags;
    DWORD dwBytesWritten = 0;

    // 初始化COM
    hr = CoInitializeEx(NULL, COINIT_MULTITHREADED);
    if (FAILED(hr)) {
        std::cerr << "CoInitializeEx failed: " << hr << std::endl;
        return hr;
    }

    // 获取设备枚举器
    hr = CoCreateInstance(
        CLSID_MMDeviceEnumerator, NULL,
        CLSCTX_ALL, IID_IMMDeviceEnumerator,
        (void**)&pEnumerator);
    if (FAILED(hr)) {
        std::cerr << "CoCreateInstance failed: " << hr << std::endl;
        // goto exit;
    }

    // 获取默认音频渲染设备
    hr = pEnumerator->GetDefaultAudioEndpoint(
        eRender, eConsole, &pDevice);
    if (FAILED(hr)) {
        std::cerr << "GetDefaultAudioEndpoint failed: " << hr << std::endl;
        // goto exit;
    }

    // 激活音频客户端
    hr = pDevice->Activate(
        IID_IAudioClient, CLSCTX_ALL,
        NULL, (void**)&pAudioClient);
    if (FAILED(hr)) {
        std::cerr << "IAudioClient activation failed: " << hr << std::endl;
        // goto exit;
    }
    // 获取音频格式
    hr = pAudioClient->GetMixFormat(&pwfx);
    if (FAILED(hr)) {
        std::cerr << "GetMixFormat failed: " << hr << std::endl;
        //goto Exit;
    }

    // 如果是WAVE_FORMAT_EXTENSIBLE（65534），则转换为标准PCM
    if (pwfx->wFormatTag == WAVE_FORMAT_EXTENSIBLE) {
        WAVEFORMATEXTENSIBLE* pExtensible = (WAVEFORMATEXTENSIBLE*)pwfx;
        if (IsEqualGUID(pExtensible->SubFormat, KSDATAFORMAT_SUBTYPE_PCM)) {
            // 如果是PCM扩展格式，降级为标准PCM
            WAVEFORMATEX* pNewFormat = (WAVEFORMATEX*)CoTaskMemAlloc(sizeof(WAVEFORMATEX));
            memcpy(pNewFormat, pwfx, sizeof(WAVEFORMATEX));
            pNewFormat->wFormatTag = WAVE_FORMAT_PCM; // 改为标准PCM格式
            CoTaskMemFree(pwfx);
            pwfx = pNewFormat;
        }
        else {
            // 如果是非PCM格式（如IEEE_FLOAT），可能需要转换
            std::cerr << "Non-PCM format detected. Consider converting to PCM." << std::endl;
        }
    }


    // 初始化音频客户端
    hr = pAudioClient->Initialize(
        AUDCLNT_SHAREMODE_SHARED,
        AUDCLNT_STREAMFLAGS_LOOPBACK,
        hnsRequestedDuration,
        0,
        pwfx,
        NULL);
    if (FAILED(hr)) {
        std::cerr << "IAudioClient Initialize failed: " << hr << std::endl;
        // goto exit;
    }

    // 获取缓冲区大小
    hr = pAudioClient->GetBufferSize(&bufferFrameCount);
    if (FAILED(hr)) {
        std::cerr << "GetBufferSize failed: " << hr << std::endl;
        // goto exit;
    }

    // 获取捕获客户端
    hr = pAudioClient->GetService(
        IID_IAudioCaptureClient,
        (void**)&pCaptureClient);
    if (FAILED(hr)) {
        std::cerr << "GetService failed: " << hr << std::endl;
        // goto exit;
    }

    // 计算实际持续时间
    hnsActualDuration = (double)REFTIMES_PER_SEC *
                        bufferFrameCount / pwfx->nSamplesPerSec;

    // 创建输出文件
    std::ofstream outFile(outputFile, std::ios::binary);
    if (!outFile.is_open()) {
        std::cerr << "Failed to open output file" << std::endl;
        hr = E_FAIL;
        // goto exit;
    }
    char* buffer = new char[1024 * 1024]; // 1MB缓冲
    outFile.rdbuf()->pubsetbuf(buffer, 1024 * 1024);

    // 写入标准WAV文件头
    DWORD dwChunkSize;
    DWORD dwRIFFSize = 36; // 初始大小(不包括data部分)
    DWORD dwFileSize;
    
    // RIFF块
    outFile.write("RIFF", 4);
    outFile.write((char*)&dwRIFFSize, 4);
    outFile.write("WAVE", 4);

    // fmt子块
    outFile.write("fmt ", 4);
    dwChunkSize = 16; // PCM格式固定16字节
    outFile.write((char*)&dwChunkSize, 4);
    
    // 写入标准PCM格式
    WORD wFormatTag = 1; // WAVE_FORMAT_PCM
    WORD nChannels = pwfx->nChannels;
    DWORD nSamplesPerSec = pwfx->nSamplesPerSec;
    DWORD nAvgBytesPerSec = pwfx->nAvgBytesPerSec;
    WORD nBlockAlign = pwfx->nBlockAlign;
    WORD wBitsPerSample = pwfx->wBitsPerSample;

    
    outFile.write((char*)&wFormatTag, 2);
    outFile.write((char*)&nChannels, 2);
    outFile.write((char*)&nSamplesPerSec, 4);
    outFile.write((char*)&nAvgBytesPerSec, 4);
    outFile.write((char*)&nBlockAlign, 2);
    outFile.write((char*)&wBitsPerSample, 2);

    // data子块
    outFile.write("data", 4);
    dwChunkSize = 0; // 初始为0
    outFile.write((char*)&dwChunkSize, 4);

    // 开始捕获
    hr = pAudioClient->Start();
    if (FAILED(hr)) {
        std::cerr << "Start failed: " << hr << std::endl;
        // goto exit;
    }

    std::cout << "Recording... Press Enter to stop" << std::endl;

    // 捕获循环
    while (true) {
        // 等待数据可用
        Sleep(hnsActualDuration/REFTIMES_PER_MILLISEC/2);

        hr = pCaptureClient->GetNextPacketSize(&packetLength);
        if (FAILED(hr)) {
            std::cerr << "GetNextPacketSize failed: " << hr << std::endl;
            break;
        }

        while (packetLength != 0) {
            // 获取捕获数据
            hr = pCaptureClient->GetBuffer(
                &pData,
                &numFramesAvailable,
                &flags, NULL, NULL);
            if (FAILED(hr)) {
                std::cerr << "GetBuffer failed: " << hr << std::endl;
                break;
            }

            if (flags & AUDCLNT_BUFFERFLAGS_SILENT) {
                pData = NULL;
            }

            // 写入文件
            if (pData && numFramesAvailable > 0) {
                DWORD dwBytesToWrite = 0;
                if (flags & AUDCLNT_BUFFERFLAGS_SILENT) {
                    // 写入静音数据(0)
                    dwBytesToWrite = numFramesAvailable * pwfx->nBlockAlign;
                    std::vector<BYTE> silence(dwBytesToWrite, 0);
                    outFile.write((char*)silence.data(), dwBytesToWrite);
                }
                else {
                    // 写入实际数据
                    dwBytesToWrite = numFramesAvailable * pwfx->nBlockAlign;
                    outFile.write((char*)pData, dwBytesToWrite);
                }
                dwBytesWritten += dwBytesToWrite;
            }


            hr = pCaptureClient->ReleaseBuffer(numFramesAvailable);
            if (FAILED(hr)) {
                std::cerr << "ReleaseBuffer failed: " << hr << std::endl;
                break;
            }

            hr = pCaptureClient->GetNextPacketSize(&packetLength);
            if (FAILED(hr)) {
                std::cerr << "GetNextPacketSize failed: " << hr << std::endl;
                break;
            }
        }

        // 检查是否有键盘输入(Enter键)
        if (GetAsyncKeyState(VK_RETURN) & 0x8000) {
            break;
        }
    }

    // 更新WAV文件头
    dwFileSize = dwBytesWritten + 36; // 36 = 头部固定大小
    dwRIFFSize = dwFileSize - 8;
    outFile.seekp(4, std::ios::beg);
    outFile.write((char*)&dwRIFFSize, 4);
    outFile.seekp(40, std::ios::beg); // data大小位置
    outFile.write((char*)&dwBytesWritten, 4);

Exit:
    // 清理资源
    if (pCaptureClient) pCaptureClient->Release();
    if (pAudioClient) pAudioClient->Release();
    if (pDevice) pDevice->Release();
    if (pEnumerator) pEnumerator->Release();
    if (pwfx) CoTaskMemFree(pwfx);
    CoUninitialize();

    return hr;
}

int main() {
    std::cout << "Audio Capture Program" << std::endl;
    std::cout << "Output file will be saved as 'output.wav'" << std::endl;
    
    HRESULT hr = RecordAudioStream(L"output.wav");
    if (FAILED(hr)) {
        std::cerr << "Recording failed with error: " << hr << std::endl;
        return 1;
    }

    std::cout << "Recording completed successfully" << std::endl;
    return 0;
}
